package bus

import chisel3._
import chisel3.util._
import common.Constants._
import bus._

class AXI_Bridge(nWords: Int) extends Module {
  val io = IO(new Bundle(){
    val in = Flipped(new CoreLinkIO(nWords*XLEN))
    val out = new AXI4
  })
  io := DontCare

  val idle :: memAddrReq :: memWriteReq :: memReadResp :: memWriteB :: release :: Nil = Enum(6)
  val reg_state = RegInit(idle)

  val dataReadBuffer = RegInit(0.U((nWords*XLEN).W))
  val dataWriterBuffer = RegInit(0.U((nWords*XLEN).W))

  val readBeatCnt = Counter(nWords)
  val writeBeatCnt = Counter(nWords)

  val startAddr = RegInit(0.U(AddrBits.W))
  val size = RegInit(0.U(2.W))
  val axi_wen = RegInit(false.B)

  // val shift = (addr(2, 0) << 3)
  // val shift_wdata = (wdata << shift).asUInt()
  io.in.req.ready := reg_state === idle
  io.in.resp.valid := (reg_state === release)
  io.in.resp.bits.rdata := Mux(reg_state === release, dataReadBuffer, 0.U)
  io.in.resp.bits.cmd := Mux(axi_wen, LinkBusCmd.writeResp, LinkBusCmd.readLast)

  val wlast = Mux(writeBeatCnt.value === (nWords - 1).U, true.B, false.B)
  val rlast = io.out.r.bits.last
  val rWord = io.out.r.bits.data
  val len = (nWords - 1).U
  val addr = Cat(io.in.req.bits.addr(AddrBits - 1, 3), 0.U(3.W))
  val wblock = WireInit(dataWriterBuffer.asTypeOf(Vec(nWords, UInt(XLEN.W))))

  switch(reg_state){
    is(idle){
      when(io.in.req.fire()){
        reg_state := memAddrReq
        size := io.in.req.bits.size
        startAddr := addr
        dataWriterBuffer := io.in.req.bits.wdata
        axi_wen := io.in.req.bits.cmd === LinkBusCmd.write
      }
    }
    is(memAddrReq){
      when(io.out.ar.fire()){
        reg_state := memReadResp
        readBeatCnt.value := (startAddr >> size) & len
      }.elsewhen(io.out.aw.fire()) {
        reg_state := memWriteReq
      }
    }
    is(memWriteReq){
      when(io.out.w.fire()){
        writeBeatCnt.inc()
        when(wlast){
          writeBeatCnt.value := 0.U
          reg_state := memWriteB
        }
      }
    }
    is(memReadResp){
      when(io.out.r.fire()){
        val nextBlock = dataReadBuffer.asTypeOf(Vec(nWords, UInt(XLEN.W)))
        nextBlock(readBeatCnt.value) := rWord

        dataReadBuffer := nextBlock.asUInt()
        readBeatCnt.inc() // a step
        when(rlast){
          readBeatCnt.value := 0.U
          reg_state := release
        }
        when (readBeatCnt.value === len) {
          readBeatCnt.value := 0.U
        }
      }
    }
    is(memWriteB){
      when(io.out.b.fire()){
        reg_state := release
      }
    }
    is(release){
      when(io.in.resp.fire()){
        reg_state := idle
      }
    }
  }

  // read addr
  io.out.ar.bits.id := 0.U
  io.out.ar.bits.addr := startAddr
  io.out.ar.bits.len := (nWords-1).U // // 协议中的 AxLen 信号从零开始表示，实际的长度值为 AxLen + 1
  io.out.ar.bits.size := size
  io.out.ar.bits.burst := AXI4Parameters.BURST_WRAP //BURST_FIXED  // 0b00, INCR
//    io.out.ar.bits.lock := 0.U
//    io.out.ar.bits.cache := 0.U
//    io.out.ar.bits.prot := 0.U
//    io.out.ar.bits.qos := 0.U
//    io.out.ar.bits.region := 0.U
//    io.out.ar.bits.user := 0.U
  io.out.ar.valid := reg_state === memAddrReq && !axi_wen

  // read data
  io.out.r.ready := reg_state === memReadResp


  // addr write
  io.out.aw.bits.id := 0.U
  io.out.aw.bits.addr := startAddr
  io.out.aw.bits.len := (nWords-1).U//0.U
  io.out.aw.bits.size := size
  io.out.aw.bits.burst := AXI4Parameters.BURST_WRAP//BURST_FIXED
//    io.out.aw.bits.lock := 0.U
//    io.out.aw.bits.cache := 0.U
//    io.out.aw.bits.prot := 0.U
//    io.out.aw.bits.qos := 0.U
//    io.out.aw.bits.region := 0.U
//    io.out.aw.bits.user := 0.U
  io.out.aw.valid := reg_state === memAddrReq && axi_wen

  // write data
  io.out.w.bits.data := wblock(writeBeatCnt.value)
  io.out.w.bits.strb := 0xff.U//strb
  io.out.w.bits.last := wlast
  io.out.w.valid := reg_state === memWriteReq && axi_wen// && !wdata_rcv

  // b
  io.out.b.ready := reg_state === memWriteB
}

// val strb =  Mux(size === 0.U,  0x1.U << startAddr(2, 0),
//             Mux(size === 1.U,  0x3.U << startAddr(2, 0),
//             Mux(size === 2.U,  0xf.U << startAddr(2, 0),
//                 0xff.U)))